import cv2
from tkinter import Label
from PIL import Image, ImageTk
from ultralytics import YOLO
import torch
class VideoPlayer:
    def __init__(self, parent):

        # 创建一个标签来显示视频
        self.label = Label(parent)
        self.label.pack()

        # 打开视频源
        #self.video_source = video_source
        self.video_source = "data/13.jpg"
        self.vid = cv2.VideoCapture(self.video_source)

        self.model = YOLO('weights/best.pt')
        self.model.to(torch.device('cpu'))
        # 开始播放视频
        self.update()

    def update(self):
        # 从视频源读取一帧
        ret, frame = self.vid.read()

        if ret:
            results = self.model(frame)
            annotated_image = results[0].plot()
            img = Image.fromarray(annotated_image)
            # 将PIL图像转换为ImageTk格式
            imgtk = ImageTk.PhotoImage(image=img)
            # 更新标签以显示新帧
            self.label.imgtk = imgtk
            self.label.configure(image=imgtk)

        # 每隔20毫秒调用一次update方法
        self.label.after(100, self.update)

    def on_closing(self):
        # 释放视频源并关闭窗口
        if self.vid.isOpened():
            self.vid.release()
